package go_pool

import (
	"errors"
	"sync"
)

// WorkerPool 定义workerPool结构体
type WorkerPool struct {
	Closed      bool   // 协程池是否关闭
	MaxNum      uint16 // 协程池最大协程数
	AliveNum    uint16 // 当前存活的协程数
	WorkerNum   uint16 // 执行协程的worker数
	mux         sync.RWMutex
	CloseOnce   sync.Once
	WorkerQueue chan *Worker           // worker的工作队列
	onPanic     func(data interface{}) // 回调函数
}

// NewWorkerPool 创建协程池：最大协程数量，当前可执行协程数量
func NewWorkerPool(maxNum, workerNum uint16) *WorkerPool {
	// worker数不能超过 最大协程数
	if workerNum > maxNum {
		workerNum = maxNum
	}
	return &WorkerPool{
		MaxNum:      maxNum,
		WorkerNum:   workerNum,
		mux:         sync.RWMutex{},
		CloseOnce:   sync.Once{},
		WorkerQueue: make(chan *Worker, maxNum),
	}
}

func (wp *WorkerPool) OnPanic(p func(data interface{})) {
	wp.onPanic = p
}

// Len 获取协程数
func (wp *WorkerPool) Len() uint16 {
	wp.mux.RLock()
	defer wp.mux.RUnlock()
	return wp.AliveNum
}

// Pause 暂停
func (wp *WorkerPool) Pause() {
	// 使所有协程阻塞。或者调整当前worker为0
	wp.AdjustSize(0)
}

// Continue 继续
func (wp *WorkerPool) Continue(num uint16) {
	// 唤醒所有协程，或者调整当前worker为暂停前的数量
	wp.AdjustSize(num)
}

// AdjustSize 调整当前协程池的数量
func (wp *WorkerPool) AdjustSize(num uint16) {
	wp.mux.Lock()
	defer wp.mux.Unlock()
	// 当前协程池已关闭，不能调整
	if wp.Closed {
		return
	}

	// 不能超过最大协程池数量
	if num > wp.MaxNum {
		num = wp.MaxNum
	}

	//判断是增加还是减少
	oldNum := wp.WorkerNum
	wp.WorkerNum = num
	// 减少worker
	if oldNum > num {
		for num < wp.AliveNum {
			wp.AliveNum--
			worker := <-wp.WorkerQueue
			worker.close()
		}
	} else {
		for num > wp.AliveNum {
			wp.AliveNum++
			worker := newWorker()
			worker.run(wp.WorkerQueue, wp.onPanic)
			wp.WorkerQueue <- worker
		}
	}
}

// Accept 接受任务
func (wp *WorkerPool) Accept(job Job) error {
	if job == nil {
		return errors.New("job cannot be nil")
	}

	wp.mux.Lock()
	select {
	case worker := <-wp.WorkerQueue:
		// 有闲置的worker
		wp.mux.Unlock()
		if worker != nil {
			worker.jobQueue <- job
		} else {
			return errors.New("workerPool is closed")
		}
	default:
		var worker *Worker
		if wp.AliveNum < wp.WorkerNum {
			// 还未达到Worker的上限
			wp.AliveNum++
			wp.mux.Unlock()
			worker = newWorker()
			worker.run(wp.WorkerQueue, wp.onPanic)
		} else if wp.AliveNum == wp.WorkerNum {
			wp.mux.Unlock()
			// 这里可能会阻塞，会阻塞的情况为：当前所有的worker都还没执行完任务就又来了新的协程任务
			worker = <-wp.WorkerQueue
		} else if wp.AliveNum > wp.WorkerNum {
			return errors.New("worker number less than alive number")
		}
		// wp.AliveNum == wp.WorkerNum 可能会出现worker = nil
		if worker != nil {
			worker.jobQueue <- job
		} else {
			return errors.New("worker pool has been closed")
		}
	}
	return nil
}

func (wp *WorkerPool) Close() {
	wp.CloseOnce.Do(func() {
		wp.mux.Lock()
		wp.Closed = true
		wp.WorkerNum = 0
		wp.MaxNum = 0
		for 0 < wp.AliveNum {
			worker := <-wp.WorkerQueue
			worker.close()
			wp.AliveNum--
		}
		close(wp.WorkerQueue)
		wp.mux.Unlock()
	})
}
